package com.lingjtx.common.thread.mq;

import com.lingjtx.common.thread.core.PoolNameConst;
import com.lingjtx.common.thread.core.ThreadPoolManager;
import com.lingjtx.common.thread.pooling.ObjectPoolTemplate;
import com.rabbitmq.stream.Consumer;
import com.rabbitmq.stream.Environment;
import com.rabbitmq.stream.Producer;
import jakarta.annotation.PreDestroy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.retry.backoff.FixedBackOffPolicy;
import org.springframework.retry.policy.SimpleRetryPolicy;
import org.springframework.retry.support.RetryTemplate;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;

public class RabbitStreamConsumer {
    private static final Logger log = LoggerFactory.getLogger(RabbitStreamConsumer.class);

    private final ObjectPoolTemplate<Environment> environmentObjectPoolTemplate;
    private final FastJsonAutoTypeMessageConverter converter;
    private final ThreadPoolManager threadPoolManager;
    private final RetryTemplate retryTemplate;

    private final Map<String, Consumer> activeConsumers = new ConcurrentHashMap<>();
    private final Semaphore rateLimiter;

    private final String deadLetterPrefix;
    private final int maxRatePerSecond;

    public RabbitStreamConsumer(ObjectPoolTemplate<Environment> environmentObjectPoolTemplate,
                                FastJsonAutoTypeMessageConverter converter,
                                ThreadPoolManager threadPoolManager,
                                RabbitStreamConsumerProperties properties) {
        this.environmentObjectPoolTemplate = environmentObjectPoolTemplate;
        this.converter = converter;
        this.threadPoolManager = threadPoolManager;
        this.deadLetterPrefix = properties.getDeadLetterPrefix();
        this.maxRatePerSecond = properties.getMaxRatePerSecond();
        this.rateLimiter = new Semaphore(maxRatePerSecond);
        this.retryTemplate = buildRetryTemplate(properties.getMaxRetries());
        startRateLimiterRefresher();
    }

    public <T> void consume(String streamName, java.util.function.Consumer<T> handler) {
        Consumer consumer = environmentObjectPoolTemplate.execute(environment -> environment.consumerBuilder()
                .stream(streamName)
                .messageHandler((context, message) -> threadPoolManager
                        .getOrCreate(PoolNameConst.RABBIT_STREAM_POOL)
                        .execute(() -> {
                            try {
                                rateLimiter.acquire();

                                byte[] body = message.getBodyAsBinary();
                                MessageProperties props = new MessageProperties();
                                Message amqpMessage = new Message(body, props);

                                retryTemplate.execute(ctx -> {
                                    @SuppressWarnings("unchecked")
                                    T obj = (T) converter.fromMessage(amqpMessage);
                                    handler.accept(obj);
                                    return null;
                                }, ctx -> {
                                    log.error("⛔️ 消息重试失败，转入死信流: {}", streamName);
                                    sendToDeadLetterStream(streamName, body);
                                    return null;
                                });

                            } catch (Exception e) {
                                log.error("❌ 消息处理失败 (非重试异常)", e);
                                sendToDeadLetterStream(streamName, message.getBodyAsBinary());
                            }
                        }))
                .build());

        activeConsumers.put(streamName, consumer);
        log.info("🚀 消费启动: {}", streamName);
    }

    private RetryTemplate buildRetryTemplate(int maxAttempts) {
        RetryTemplate template = new RetryTemplate();

        SimpleRetryPolicy retryPolicy = new SimpleRetryPolicy(maxAttempts);
        FixedBackOffPolicy backOff = new FixedBackOffPolicy();
        backOff.setBackOffPeriod(500); // 500ms
        template.setRetryPolicy(retryPolicy);
        template.setBackOffPolicy(backOff);

        return template;
    }

    private void sendToDeadLetterStream(String originalStream, byte[] data) {
        String deadLetterStream = deadLetterPrefix + originalStream;
        Producer producer = environmentObjectPoolTemplate.execute(environment -> environment.producerBuilder().stream(deadLetterStream).build());
        producer.send(producer.messageBuilder().addData(data).build(), confirmationStatus -> {
            if (!confirmationStatus.isConfirmed()) {
                log.error("Message failed to reach stream: {}", originalStream);
            }
        });
        producer.close();
    }

    private void startRateLimiterRefresher() {
        ScheduledExecutorService scheduler = threadPoolManager.getOrCreateScheduled(PoolNameConst.RABBIT_STREAM_SCHEDULER);
        scheduler.scheduleAtFixedRate(() -> {
            int permitsToAdd = maxRatePerSecond - rateLimiter.availablePermits();
            rateLimiter.release(permitsToAdd);
        }, 1, 1, TimeUnit.SECONDS);
    }

    @PreDestroy
    public void shutdown() {
        activeConsumers.values().forEach(consumer -> {
            try {
                consumer.close();
            } catch (Exception e) {
                log.warn("关闭 consumer 失败", e);
            }
        });
        threadPoolManager.shutdown(PoolNameConst.RABBIT_STREAM_POOL);
        threadPoolManager.shutdown(PoolNameConst.RABBIT_STREAM_SCHEDULER);
        log.info("🛑 所有 RabbitMQ Stream 消费者已关闭");
    }
}

